#!/usr/bin/python
'''
Takes a GFF file, followed by fasta files and splits sam up by fasta names.
'''
import sys
import os

# parameters.
gff_file = sys.argv[1]

fasta_files = []
for i in range(2, len(sys.argv)):
	fasta_files.append(sys.argv[i])
	
# index fasta files.
def index_fasta(f, fout):
	# open file.
	fin = open(f, "rb")
	lines = fin.readlines()
	fin.close()
	
	# pull out names.
	names = {}
	for line in lines:
		if line[0] != ">": continue
		names[line.strip()[1::]] = fout
	
	del lines
	return names
	
# Index fasta.
flist = []
fouts = {}
for f in fasta_files:
	# make basename.
	bname1 = os.path.basename(f)
	bname2 = os.path.basename(gff_file)
	
	# open output file.
	fout = open("%s_%s" % (bname1, bname2), "wb")
	flist.append(fout)
	
	# open file.
	fin = open(f, "rb")
	lines = fin.readlines()
	fin.close()
		
	# pull out names.
	for line in lines:
		if line[0] != ">": continue
		fouts[line.strip()[1::]] = fout
	
	del lines

	
# open gff file.
fin = open(gff_file, "rb")

# loop over file in chunks.
i = 0
sz = 0
chunk = 0
while 1 == 1:
	# check if we need new chunk
	if i == sz:
		# read new chunk.
		lines = fin.readlines(100000000)
		sz = len(lines)
		i = 0
		chunk += 1
		print chunk, sz
		
	# check if we are done.
	if sz == 0: break
	
	# sanity.
	tmp = lines[i].split("\t")
	if tmp[2] not in fouts: 
		i += 1
		continue
		
	# look up file.
	fouts[tmp[2]].write(lines[i])
	i += 1
	
# close files.
for f in flist:
	f.close()
